# model_wrappers.py

import onnxruntime as ort
from openvino.runtime import Tensor
import cv2
import numpy as np
import time


class ONNXModelWrapper:
    def __init__(self, session):
        self.session = session
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]
        self.input_shape = self.session.get_inputs()[0].shape  # 通常是 [1, 3, 640, 640]

    def preprocess(self, image):
        """预处理图像到模型输入格式"""
        t0 = time.time()
        input_h, input_w = self.input_shape[2], self.input_shape[3]
        resized = cv2.resize(image, (input_w, input_h))
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
        chw = np.transpose(rgb, (2, 0, 1))
        normalized = chw.astype(np.float32) / 255.0
        batched = np.expand_dims(normalized, axis=0)
        t1 = time.time()
        return batched, (t1 - t0) * 1000  # 返回预处理时间（毫秒）

    def postprocess(self, outputs, orig_shape, t0):
        """后处理模型输出，返回 boxes, confs, class_ids"""
        t_start = time.time()
        pred = outputs[0][0]  # [84, 8400]
        pred = np.transpose(pred, (1, 0))  # [8400, 84]

        boxes = pred[:, :4]  # cx, cy, w, h
        scores = pred[:, 4:]  # class scores

        height, width = orig_shape[:2]
        input_h, input_w = self.input_shape[2], self.input_shape[3]

        x1 = (boxes[:, 0] - boxes[:, 2] / 2) * width / input_w
        y1 = (boxes[:, 1] - boxes[:, 3] / 2) * height / input_h
        x2 = (boxes[:, 0] + boxes[:, 2] / 2) * width / input_w
        y2 = (boxes[:, 1] + boxes[:, 3] / 2) * height / input_h

        boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)

        confs = np.max(scores, axis=1)
        cls_ids = np.argmax(scores, axis=1)

        mask = confs > 0.25
        boxes_xyxy = boxes_xyxy[mask]
        confs = confs[mask]
        cls_ids = cls_ids[mask]

        if len(boxes_xyxy) > 0:
            indices = cv2.dnn.NMSBoxes(boxes_xyxy.tolist(), confs.tolist(), 0.25, 0.45)
            if len(indices) == 0:
                boxes_xyxy = np.array([])
                confs = np.array([])
                cls_ids = np.array([])
            else:
                if isinstance(indices, np.ndarray):
                    indices = indices.flatten()
                boxes_xyxy = boxes_xyxy[indices]
                confs = confs[indices]
                cls_ids = cls_ids[indices]

        names = {i: f"class{i}" for i in range(scores.shape[1])}

        t_end = time.time()
        postprocess_ms = (t_end - t_start) * 1000
        return [{
            'boxes': boxes_xyxy,
            'confs': confs,
            'cls_ids': cls_ids,
            'names': names
        }], postprocess_ms

    def predict_with_timing(self, image):
        t0 = time.time()
        input_tensor, preprocess_ms = self.preprocess(image)
        t1 = time.time()

        outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
        t2 = time.time()

        result, postprocess_ms = self.postprocess(outputs, image.shape, t0)
        t3 = time.time()

        inference_ms = (t2 - t1) * 1000

        return result, preprocess_ms, inference_ms, postprocess_ms


class OpenVINOModelWrapper:
    def __init__(self, compiled_model, class_names=None):
        self.compiled_model = compiled_model
        self.input_layer = self.compiled_model.inputs[0]
        self.output_layer = self.compiled_model.outputs[0]
        self.infer_request = self.compiled_model.create_infer_request()
        self.class_names = class_names or []  # 👈 保存

    def preprocess(self, image):
        t0 = time.time()
        _, _, h, w = self.input_layer.shape
        resized = cv2.resize(image, (w, h))
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
        chw = np.transpose(rgb, (2, 0, 1))
        normalized = chw.astype(np.float32) / 255.0
        batched = np.expand_dims(normalized, axis=0)
        t1 = time.time()
        return batched, (t1 - t0) * 1000

    def postprocess(self, outputs, orig_shape, t0, meta=None):
        t_start = time.time()
        try:
            pred = outputs[0]

            # 处理 [1, 8400, C] 或 [8400, C]
            if len(pred.shape) == 3:
                if pred.shape[0] == 1:
                    pred = pred[0]  # (8400, C)
                else:
                    pred = pred[0]  # 取 batch 0

            num_classes = pred.shape[1] - 4
            if num_classes <= 0:
                print(f"[ERROR] Invalid number of classes: {num_classes}")
                return [], 0

            # 解码归一化坐标 (值在 [0,1])
            input_size = 640.0
            cx = pred[:, 0] * input_size
            cy = pred[:, 1] * input_size
            w = pred[:, 2] * input_size
            h = pred[:, 3] * input_size

            x1 = cx - w / 2
            y1 = cy - h / 2
            x2 = cx + w / 2
            y2 = cy + h / 2
            boxes_640 = np.stack([x1, y1, x2, y2], axis=1)

            scores = pred[:, 4:]
            confs = np.max(scores, axis=1)
            cls_ids = np.argmax(scores, axis=1)

            # NMS
            indices = cv2.dnn.NMSBoxes(boxes_640.tolist(), confs.tolist(), 0.25, 0.45)
            if len(indices) == 0:
                return [], (time.time() - t_start) * 1000

            indices = indices.flatten()
            boxes_640 = boxes_640[indices]
            confs = confs[indices]
            cls_ids = cls_ids[indices]

            # ========== 映射回原始图像 ==========
            if meta is None:
                boxes_original = boxes_640
            else:
                top_pad, left_pad = meta['pad']
                scale = meta['scale']

                boxes_resized = boxes_640.copy()
                boxes_resized[:, [0, 2]] -= left_pad
                boxes_resized[:, [1, 3]] -= top_pad
                boxes_resized[:, [0, 2]] = np.clip(boxes_resized[:, [0, 2]], 0, meta['resized_shape'][1])
                boxes_resized[:, [1, 3]] = np.clip(boxes_resized[:, [1, 3]], 0, meta['resized_shape'][0])

                boxes_original = boxes_resized / scale

            names = {i: self.class_names[i] if i < len(self.class_names) else f"class{i}"
                     for i in range(num_classes)}

            result = [{
                'boxes': boxes_original,
                'confs': confs,
                'cls_ids': cls_ids,
                'names': names
            }]

            postprocess_ms = (time.time() - t_start) * 1000
            return result, postprocess_ms

        except Exception as e:
            print(f"[ERROR] Postprocessing failed: {e}")
            import traceback
            traceback.print_exc()
            return [], 0

    def predict_with_timing(self, image, meta=None):
        t0 = time.time()
        input_tensor, preprocess_ms = self.preprocess(image)
        t1 = time.time()

        self.infer_request.infer({self.input_layer.any_name: input_tensor})
        outputs = self.infer_request.get_output_tensor().data
        t2 = time.time()

        result, postprocess_ms = self.postprocess([outputs], image.shape, t0, meta=meta)
        inference_ms = (t2 - t1) * 1000

        return result, preprocess_ms, inference_ms, postprocess_ms